package model;

import com.example.entity.University;
import com.example.mapper.UniversityRepository;
import org.junit.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.boot.test.context.SpringBootTest;
import weka.core.DenseInstance;
import weka.core.Instances;
import weka.experiment.InstanceQuery;
import weka.classifiers.functions.SMOreg;
import weka.classifiers.Evaluation;

import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.List;
import weka.core.Attribute;

@Component
public class ModelTraining {

    @Autowired
    private UniversityRepository universityRepository;

    @Test
    public void trainAndEvaluateModel() throws Exception {
        // 从数据库获取特定学校和专业的历史数据
        List<University> universities = universityRepository.findAll();

        // 准备用于模型训练的数据集
        Instances dataset = prepareDataset(universities);

        // 网格搜索
        double[] C_values = {0.1, 0.5, 1.0, 2.0};
        double bestC = 0;
        double bestMSE = Double.MAX_VALUE;
        SMOreg bestModel = null;

        for (double C_value : C_values) {
            SMOreg smoreg = new SMOreg();
            smoreg.setC(C_value);
            smoreg.buildClassifier(dataset);

            Evaluation eval = new Evaluation(dataset);
            eval.crossValidateModel(smoreg, dataset, 10, new java.util.Random(1));

            double mse = eval.rootMeanSquaredError();
            if (mse < bestMSE) {
                bestMSE = mse;
                bestC = C_value;
                bestModel = smoreg;
            }
        }

        System.out.println("Best C: " + bestC);
        System.out.println("Best MSE: " + bestMSE);
    }

    private Instances prepareDataset(List<University> universities) {
        ArrayList<Attribute> attributes = new ArrayList<>();
        attributes.add(new Attribute("bbl")); // 报录比
        attributes.add(new Attribute("number")); // 招生人数
        attributes.add(new Attribute("methodNumeric")); // 考试形式的数值表示（统考=1，自命题=0）
        attributes.add(new Attribute("line")); // 分数线

        Instances dataset = new Instances("UniversityData", attributes, universities.size());
        for (University uni : universities) {
            double[] vals = new double[dataset.numAttributes()];
            vals[0] = uni.getBbl();
            vals[1] = uni.getNumber();
            vals[2] = "统考".equals(uni.getMethod()) ? 1.0 : 0.0;
            vals[3] = uni.getLine();

            dataset.add(new DenseInstance(1.0, vals));
        }
        dataset.setClassIndex(dataset.numAttributes() - 1);
        return dataset;
    }

    private void performGridSearch(Instances dataset) throws Exception {
        double[] C_values = {0.1, 0.5, 1.0, 2.0}; // 正则化参数C的范围
        double bestC = 0;
        double bestMSE = Double.MAX_VALUE;
        SMOreg bestModel = null;

        for (double C_value : C_values) {
            SMOreg smoreg = new SMOreg();
            smoreg.setC(C_value);
            smoreg.buildClassifier(dataset); // 训练模型

            // 使用交叉验证评估模型
            Evaluation eval = new Evaluation(dataset);
            eval.crossValidateModel(smoreg, dataset, 10, new java.util.Random(1));

            // 记录最佳参数和模型
            double mse = eval.rootMeanSquaredError();
            if (mse < bestMSE) {
                bestMSE = mse;
                bestC = C_value;
                bestModel = smoreg;
            }
        }

        // 模型评估结果输出
        System.out.println("Best C: " + bestC);
        System.out.println("Best MSE: " + bestMSE);
    }
}

